import matplotlib.pyplot as plt

decision_node = dict(boxstyle='sawtooth', fc='0.8')
leaf_node = dict(boxstyle='round4', fc='0.8')
arrow_args = dict(arrowstyle='<-')


def plot_node(node_text, center, parent, node_type):
    create_plot.ax1.annotate(node_text, xy=parent, xycoords='axes fraction', xytext=center, textcoords='axes fraction',
                             va='center', ha='center', bbox=node_type, arrowprops=arrow_args)


def get_num_leafs(tree):
    num_leafs = 0
    first_str = list(tree.keys())[0]
    second_dict = tree[first_str]
    for key in second_dict.keys():
        if isinstance(second_dict[key], dict):
            num_leafs += get_num_leafs(second_dict[key])
        else:
            num_leafs += 1
    return num_leafs


def get_tree_depth(tree):
    max_depth = 0
    first_str = list(tree.keys())[0]
    second_dict = tree[first_str]
    for key in second_dict.keys():
        if isinstance(second_dict[key], dict):
            this_depth = get_tree_depth(second_dict[key]) + 1
        else:
            this_depth = 1
        if this_depth > max_depth:
            max_depth = this_depth
    return max_depth


def plot_mid_text(center, parent, text_string):
    x_mid = (parent[0] - center[0]) / 2.0 + center[0]
    y_mid = (parent[1] - center[1]) / 2.0 + center[1]
    create_plot.ax1.text(x_mid, y_mid, text_string)


def plot_tree(tree, parent_center, node_text):
    num_leafs = get_num_leafs(tree)
    depth = get_tree_depth(tree)
    first_str = list(tree.keys())[0]
    center_pt = (plot_tree.x_offset + (1.0 + float(num_leafs)) / 2.0 / plot_tree.total_w,
                 plot_tree.y_offset)
    plot_mid_text(center_pt, parent_center, node_text)
    plot_node(first_str, center_pt, parent_center, decision_node)
    second_dict = tree[first_str]
    plot_tree.y_offset = plot_tree.y_offset - 1.0 / plot_tree.total_d
    for key in second_dict.keys():
        if isinstance(second_dict[key], dict):
            plot_tree(second_dict[key], center_pt, str(key))
        else:
            plot_tree.x_offset = plot_tree.x_offset + 1.0 / plot_tree.total_w
            plot_node(second_dict[key], (plot_tree.x_offset, plot_tree.y_offset), center_pt, leaf_node)
            plot_mid_text((plot_tree.x_offset, plot_tree.y_offset), center_pt, str(key))
    plot_tree.y_offset = plot_tree.y_offset + 1.0 / plot_tree.total_d


def create_plot(tree):
    fig = plt.figure(1, facecolor='white')
    fig.clf()
    ax_props = dict(xticks=[], yticks=[])
    create_plot.ax1 = plt.subplot(111, frameon=False, **ax_props)
    plot_tree.total_w = float(get_num_leafs(tree))
    plot_tree.total_d = float(get_tree_depth(tree))
    plot_tree.x_offset = -0.5 / plot_tree.total_w
    plot_tree.y_offset = 1.0
    plot_tree(tree, (0.5, 1.0), '')
    plt.show()
